package cn.lg.soar.mvc.config;

import cn.lg.soar.common.util.IpUtil;
import cn.lg.soar.common.util.data.DataUtil;
import cn.lg.soar.mvc.model.IpFilterProperties;
import cn.lg.soar.mvc.model.IpFilterRule;
import cn.lg.soar.mvc.util.ResponseUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.annotation.Order;
import org.springframework.util.CollectionUtils;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * ip拦截
 * @author luguoxiang
 */
@Order(0)
@Configuration
public class IpFilterInterceptor implements WebMvcConfigurer {

    protected static final Logger LOGGER = LoggerFactory.getLogger(IpFilterInterceptor.class);
    @Autowired
    private IpFilterProperties props;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        // 白名单
        List<IpFilterRule> whitelist = props.getWhitelist();
        if (DataUtil.isValuable(whitelist)) {
            // 循环增加拦截器
            for (IpFilterRule rule : whitelist) {
                Set<String> paths = rule.getPaths();
                Set<String> ips = rule.getIps();
                if (CollectionUtils.isEmpty(paths) || CollectionUtils.isEmpty(ips)) {
                    throw new RuntimeException("ip-whitelist configuration error：Must include the “path-patterns” and “ips”");
                }
                // 格式化IP
                Set<String> ipSet = ips.stream().map(IpUtil::fmtIp).collect(Collectors.toSet());
                // 拦截的路径
                registry.addInterceptor(new WhitelistInterceptor(ipSet, rule.getMessage()))
                        .addPathPatterns(paths.toArray(new String[0]))
                        .order(0);
            }
        }

        // 黑名单
        List<IpFilterRule> blacklist = props.getBlacklist();
        if (DataUtil.isValuable(blacklist)) {
            // 循环增加拦截器
            for (IpFilterRule rule : blacklist) {
                Set<String> paths = rule.getPaths();
                Set<String> ips = rule.getIps();
                if (CollectionUtils.isEmpty(paths) || CollectionUtils.isEmpty(ips)) {
                    throw new RuntimeException("ip-blacklist configuration error：Must include the “path-patterns” and “ips”");
                }
                // 格式化IP
                Set<String> ipSet = ips.stream().map(IpUtil::fmtIp).collect(Collectors.toSet());
                // 拦截的路径
                registry.addInterceptor(new BlacklistInterceptor(ipSet, rule.getMessage()))
                        .addPathPatterns(paths.toArray(new String[0]))
                        .order(0);
            }
        }
    }

    /**
     * 白名单拦截器
     */
    public static class WhitelistInterceptor implements HandlerInterceptor {

        private final Set<String> ips;
        private final String message;

        public WhitelistInterceptor(Set<String> ips, String message) {
            this.ips = ips;
            this.message = message;
        }

        @Override
        public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
            String remoteIp = IpUtil.getRemoteIp(request);
            if (this.ips.contains(remoteIp)) {
                return true;
            }
            LOGGER.warn("IP白名单拦截，IP={}", remoteIp);
            ResponseUtils.E412(response, message);
            return false;
        }

    }

    /**
     * 黑名单拦截器
     */
    public static class BlacklistInterceptor implements HandlerInterceptor {

        private final Set<String> ips;
        private final String message;

        public BlacklistInterceptor(Set<String> ips, String message) {
            this.ips = ips;
            this.message = message;
        }

        @Override
        public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
            String remoteIp = IpUtil.getRemoteIp(request);
            if (this.ips.contains(remoteIp)) {
                LOGGER.warn("IP黑名单拦截，IP={}", remoteIp);
                ResponseUtils.E412(response, message);
                return false;
            }
            return true;
        }

    }

}
